# -*- coding: utf-8 -*-
"""
@Time    : 2025/2/13 18:22 
@Author  : ZhangShenao 
@File    : 2.使用底层接口调用模型.py 
@Desc    : 使用底层接口调用模型
"""

from transformers import AutoTokenizer, AutoModelForCausalLM

# 创建Tokenizer分词器
# Tokenizer用于实现分词，以及文本与token之间的编解码
# 通常不同的模型会采用不同的Tokenizer,因此Tokenizer需要与模型配合使用
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-0.5B-Instruct")

# 加载模型
model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-0.5B-Instruct")

# 将文本输入转换成Token
messages = [
    {"role": "user", "content": "请写一首赞美春天的诗，要求不包含春字"},
]
text = tokenizer.apply_chat_template(
    messages,
    tokenize=False,
    add_generation_prompt=True
)
model_inputs = tokenizer([text], return_tensors="pt").to(model.device)

# 调用大模型,生成结果Token
generated_ids = model.generate(
    **model_inputs,
    max_new_tokens=512
)
generated_ids = [
    output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
]

# 将生成的结果Token解码成文本
response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
print(response)
